import numpy as np
from multiprocessing import Process, Array,Lock
from scipy import signal
import asa
from solve import solve


asa_settings={'rand_seed':905,
              #'test_in_cost_func':0, # not accessible
              #'use_rejected_cost':0, # not accessible
              'asa_out_file':'asageneexpression.log',
              'limit_acceptances':1000,
              'limit_generated':99999,
              'limit_invalid_generated_states':1000, # check
              'accepted_to_generated_ratio':0.0001,
              'cost_precision':1e-18,
              'maximum_cost_repeat':5,
              'number_cost_samples':5,
              'temperature_ratio_scale':0.001,
              'cost_parameter_scale_ratio':0.5,
              'temperature_anneal_scale':120,
              'include_integer_parameters':0,
              'user_initial_parameters':1,
              'sequential_parameters':-3,
              'initial_parameter_temperature':1,
              'acceptance_frequency_modulus':100,
              'generated_frequency_modulus':10000,
              'reanneal_cost':1,
              'reanneal_parameters':1,
              'delta_x':0.001
              }

#Flags
qpcrData=False
plotoutput = True
# Constants
windowwidth = 20

# Functions
def KL_divergence(p, q):
    """ Compute KL divergence, (used to compute J-S divergence)."""
    idx=(q>0)&(p>0)
    #print('qmin',np.min(q[idx]),'pmin',np.min(p[idx]))
    kld=np.sum(p[idx]*np.log(p[idx]/q[idx]))
    return kld

def Jensen_Shannon_divergence(p, q):
    """ Gives the Jensen-Shannon divergence. """
    weight = 0.5
    average = weight * p + (1 - weight) * q
    jsd = weight * KL_divergence(p, average) + (1 - weight) * KL_divergence(q, average)
    print(jsd)
    return jsd
    
def Jensen_Shannon_divergence_penalty(p, q):
    """ Gives the Jensen-Shannon divergence. """
    weight = 0.5
    average = weight * p + (1 - weight) * q
    jsdp = weight * KL_divergence(p, average) + (1 - weight) * KL_divergence(q, average) + 10
    print(jsdp)
    return jsdp

def cost_function(x,Ddata):
    if qpcrData:
        NM = np.round((x[4]*Ddata))
    else:
        NM=Ddata
    Min = 0.
    Max = NM.max()
    xvalbnd= 0.5+np.arange(Min-1, Max+1)
    MaxH=np.int(np.round(Max/2.)+1)
    #Select number of processors for parallel computation of steady state distributions
    NPROC=2
    Pm=Array('d', MaxH)
    proc=[]
    for i in range(NPROC):
        Ns=range(i,MaxH,NPROC)
        proc.append(Process(target=solve, args=(x[0],x[1],x[2],x[3],Ns,Pm)))
        proc[i].start()
    for i in range(NPROC):
        proc[i].join()
    p=np.array(Pm)

    D = np.convolve(p, p)
    window = signal.gaussian(windowwidth, std=7)
    q = np.convolve(D.T, window, 'same')/window.sum()

    pmoleculenumbers = np.histogram(NM, xvalbnd)[0]/np.float(NM.size)
    p = np.convolve(pmoleculenumbers,window, 'same')/window.sum()
    pqsize=np.minimum(p.size,q.size)
    p=p[:pqsize]
    q=q[:pqsize]
    
    if np.sum(p)<0.93:
       check1 = 1.
    else:
       check1 = 0.

    if np.sum(q)<0.93:
       check2 = 1.
    else:
       check2 = 0.
    
    if check2 == 1.:
       print('Generated distribution q is not accurate, adding penalty of 10 to cost. This will be corrected')

    if (check1 or check2):
       return Jensen_Shannon_divergence_penalty(p, q)
    else:
       return Jensen_Shannon_divergence(p, q)

 #   return Jensen_Shannon_divergence(p, q)

if __name__ == '__main__':
    p = [1.0/10, 1.0/10, 0]
    q = [0, 1.0/10, 9.0/10]
    Ddata=np.loadtxt('insert_Data_file_location_here')
    cost=lambda x: cost_function(x,Ddata)
    upper = 2400/(np.amax(Ddata))
    lower = 1/(np.amin(Ddata[np.nonzero(Ddata)]))
    x0 = np.array([1, 100, 5000, 1],dtype=np.float64)
    xmin = np.array([0.1, 0.1, 1, 1],dtype=np.float64)
    xmax = np.array([20, 200, 20000, 1],dtype=np.float64)
    if qpcrData:
        x0=np.r_[x0,lower+(upper-lower)/10]
        xmin=np.r_[xmin,lower]
        xmax=np.r_[xmax,upper]
    print(asa.asa(func=cost, x0=x0, xmin=xmin, xmax=xmax, full_output=True,**asa_settings))
    
#plot data and fitting
if plotoutput:True
from matplotlib import pyplot as plt
fig = plt.figure()
plt.plot(p)
plt.plot(q)
